// Copyright 2023 Huawei Cloud Computing Technology Co., Ltd.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "Heartbeat.h"
#include <ctime>
#include <sys/prctl.h>
#include <unistd.h>

#include "HWSecure/include/securec.h"
#include "Log/logging.h"
#include "MemoryManagement/VmiBuffer.h"

#if ENABLE_STATISTICS
#include "Statistics/Statistics.h"
#include "Statistics/StatisticsObject.h"
#endif

namespace {
    enum HEARTBEAT_PACKAGE_TYPE { HEARTBEAT_REQ = 0, HEARTBEAT_RES = 1};
    const unsigned int MICROSECOND_OFFSET = 500;
    const int CHECK_DURATION = 30;  // check duration
    const int AVERAGE_NUM = 10;      // calculate network average delay with the last 10 heartbeat data
    const int HEARTBEAT_SEND_INTERVAL = 100;   // send heartbeat interval
    const int SECOND_TO_MICRO = 1000000;
    const int SECOND_TO_MILL = 1000;
    const int NANO_TO_MICRO = 1000;
    const int THREAD_NAME_LEN = 32;
    const int GET_DATA_FROM_QUEUE_MIN_WAIT_TIME = 1000;
}

/**
 * @brief: get microsecond
 * @param [in] t: time
 * @return: microsecond
 */
static inline uint64_t GetMicroSecond(struct timespec t)
{
    return ((uint64_t)t.tv_sec * SECOND_TO_MICRO +
        ((uint64_t)t.tv_nsec + MICROSECOND_OFFSET) / NANO_TO_MICRO);
}

/**
 * @brief: heartbeat has for function:
 *               1、timing send heartbeat (SendHeartbeatRequest);
 *               2、handle recv heartbeat (RecvHeartbeatResponse);
 *               3、compute net loop;
 *               4、check net status.
 * @param [in] param: heartbeat param
 */
void TaskentryHeartbeat(std::shared_ptr<Heartbeat::ThreadPara> param)
{
    INFO("HEARTBEAT thread entering");

    char threadName[THREAD_NAME_LEN] = {'\0'};
    int code = snprintf_s(threadName, THREAD_NAME_LEN, sizeof(threadName), "HB-%p", param->GetVmiSocket().get());
    if (code < 0) {
        // not need return
        ERR("snprintf_s error");
    }
    // format is 0xHHHHHHHHHHHHHHHH
    prctl(PR_SET_NAME, threadName);

    StreamMsgHead* streamMessage = reinterpret_cast<StreamMsgHead *>(VmiAllocBuffer(
        sizeof(StreamMsgHead) + sizeof(HeartbeatMsg)));
    if (streamMessage != nullptr) {
        struct timespec sendTime = {0, 0};
        clock_gettime(CLOCK_MONOTONIC, &sendTime);
        param->SetNextSendTime(GetMicroSecond(sendTime));

        while (param->GetStopFlag() != true) {
            param->SendHeartbeatRequest(*streamMessage);
            param->RecvHeartbeatResponse();
            param->CheckNetStatus();
        }

        VmiFreeBuffer(streamMessage);
    } else {
        ERR("alloc streamMessage failed");
    }

    ERR("thread exited");
    return;
}

/**
 * @brief: heartbeat thread use this data
 */
/*lint -save -e1401 */
Heartbeat::ThreadPara::ThreadPara() : m_msgFragment(), m_msgReassemble(VMIMsgType::HEARTBEAT)
{
    m_stopFlag = false;
    m_sendSeqNum = 0;
    m_recvSeqNum = 0;
    m_nextSendTime = 0;
    m_netLoopbackLatencyAverage = 0;
    m_netLoopbackLatencyMax = 0;
}
/*lint -restore */
/**
 * @brief: set m_nextSendTime
 * @param [in] nextTime
 */
void Heartbeat::ThreadPara::SetNextSendTime(uint64_t nextTime)
{
    m_nextSendTime = nextTime;
}

/**
 * @brief: get m_netLoopbackLatencyAverage
 * @return: m_netLoopbackLatencyAverage
 */
uint64_t Heartbeat::ThreadPara::GetNetLoopbackLatencyAverage()
{
    return m_netLoopbackLatencyAverage;
}

/**
 * @brief: get net loop max
 * @return: loop, 0 is Infinity
 */
uint64_t Heartbeat::ThreadPara::GetNetLoopbackLatencyMax()
{
    return m_netLoopbackLatencyMax;
}

/**
 * @brief: get m_packetQueue
 * @return: m_packetQueue
 */
InterThreadComm& Heartbeat::ThreadPara::GetInterThreadComm()
{
    return m_packetQueue;
}

/**
 * @brief: get m_stopFlag
 * @return: m_stopFlag
 */
bool Heartbeat::ThreadPara::GetStopFlag()
{
    return m_stopFlag;
}

/**
 * @brief: set m_stopFlag
 * @param [in] stopFlag: stopFlag
 */
void Heartbeat::ThreadPara::SetStopFlag(bool stopFlag)
{
    m_stopFlag = stopFlag;
}

/**
 * @brief: set m_vmiSocket
 * @param [in] vmiSocket: vmiSocket
 */
void Heartbeat::ThreadPara::SetVmiSocket(std::shared_ptr<VmiSocket> vmiSocket)
{
    m_vmiSocket = vmiSocket;
}

/**
 * @brief: get m_vmiSocket
 * @return: m_vmiSocket
 */
std::shared_ptr<VmiSocket> Heartbeat::ThreadPara::GetVmiSocket()
{
    return m_vmiSocket;
}

/**
 * @brief: set m_callback
 * @param [in] callback: callback
 */
void Heartbeat::ThreadPara::SetCallback(std::function<void(std::shared_ptr<VmiSocket>)> callback)
{
    m_callback = callback;
}

/**
 * @brief: send heartbeat request
 * @param [in] streamMessage: socket package
 */
void Heartbeat::ThreadPara::SendHeartbeatRequest(StreamMsgHead& streamMessage)
{
    struct timespec sendTime = {0, 0};

    clock_gettime(CLOCK_MONOTONIC, &sendTime);
    HeartbeatMsg* heartbeatMessage = reinterpret_cast<HeartbeatMsg *>(&streamMessage + 1);
    heartbeatMessage->sendMicroSecond = GetMicroSecond(sendTime);
    if ((heartbeatMessage->sendMicroSecond + MICROSECOND_OFFSET) < m_nextSendTime) { // Rounding
        return;
    }
    if (m_latencyQueue.size() >= static_cast<size_t>(CHECK_DURATION)) {
        m_latencyQueue.pop_front();
    }

    heartbeatMessage->type = HEARTBEAT_REQ;
    heartbeatMessage->seqNum = ++m_sendSeqNum;

#if ENABLE_STATISTICS
    // statistics heartbeat
    Statistics *statistics = Statistics::GetInstance();
    if (statistics != nullptr) {
        static StatisticsObject* netStatisticsObject = statistics->AllocStatisticsObject(NET_TABLE_NAME,
                                                                                         NET_TABLE_COLUMNS,
                                                                                         SAMPLE_INTERVAL_PERSECOND);
        if (netStatisticsObject != nullptr) {
            static int indexHeartSeq = netStatisticsObject->GetComponentIndex(NET_COLUMN_HEART_SEQ);
            static int indexHeartSendTime = netStatisticsObject->GetComponentIndex(NET_COLUMN_HEARTBEAT_SEND_TIME);
            netStatisticsObject->UpdateComponent(indexHeartSeq, UPDATE, heartbeatMessage->seqNum);
            netStatisticsObject->UpdateComponent(indexHeartSendTime, UPDATE, heartbeatMessage->sendMicroSecond);
        }
    }
#endif

    uint64_t sendMicroSecond = heartbeatMessage->sendMicroSecond;
    m_latencyQueue.emplace_back(std::make_pair(sendMicroSecond, static_cast<uint64_t>(0)));
    m_nextSendTime += HEARTBEAT_SEND_INTERVAL * SECOND_TO_MILL;

    if (m_msgFragment.FragmentSend(VMIMsgType::HEARTBEAT, m_vmiSocket, streamMessage, sizeof(HeartbeatMsg))
        != (sizeof(HeartbeatMsg))) {
        ERR("Send heartbeat msg failed send seqNum: %ju", m_sendSeqNum);
    }

    return;
}

struct StreamHeartbeatMsg {
    struct StreamMsgHead header;
    struct HeartbeatMsg heartbeatMsg;
} __attribute__((packed));

/**
 * @brief: handle recv heartbeat message
 * @param [in] reassembleMsg: recv message
 * @param [in] recvTime: recvTime
 */
void Heartbeat::ThreadPara::HandleRecvMsg(const std::pair<uint8_t*, uint32_t>& reassembleMsg,
    struct timespec recvTime)
{
    uint8_t* msgData = reassembleMsg.first;
    uint32_t msgSize = reassembleMsg.second;
    auto heartbeatMsg = reinterpret_cast<HeartbeatMsg *>(msgData);
    if (heartbeatMsg->type == HEARTBEAT_REQ) {
        heartbeatMsg->type = HEARTBEAT_RES;

        StreamHeartbeatMsg loopbackMsg = { { }, { 0, 0, 0, {0} } };
        errno_t rc = memcpy_s(reinterpret_cast<uint8_t *>(&loopbackMsg) + sizeof(StreamMsgHead),
            msgSize, msgData, msgSize);
        if (rc != EOK) {
            ERR("error: memcpy failed, rc:%d", rc);
            return;
        }

        auto head = reinterpret_cast<StreamMsgHead *>(&loopbackMsg);
        if (m_msgFragment.FragmentSend(VMIMsgType::HEARTBEAT, m_vmiSocket,
            *head, sizeof(HeartbeatMsg))
            != sizeof(HeartbeatMsg)) {
            ERR("Loopback heartbeat msg failed, seqnum(%ju), recvSeqNum(%ju)", heartbeatMsg->seqNum, m_recvSeqNum);
        }
    } else if (heartbeatMsg->type == HEARTBEAT_RES) {
        if (heartbeatMsg->seqNum > m_recvSeqNum) {
            CalcLatency(*heartbeatMsg, GetMicroSecond(recvTime));
        } else {
            ERR("Heartbeat response msg seqnum(%ju) is error, recvSeqNum(%ju)", heartbeatMsg->seqNum, m_recvSeqNum);
        }
    } else {
        ERR("Heartbeat message type is error!");
    }
}

/**
 * @brief: handle recv heartbeat message
 */
void Heartbeat::ThreadPara::RecvHeartbeatResponse()
{
    struct timespec recvTime = {0, 0};
    StreamMsgHead* recvMessage = nullptr;
    std::pair<uint8_t*, uint32_t> responseMessage(nullptr, 0);
    uint64_t waitTime = 0;

    do {
        clock_gettime(CLOCK_MONOTONIC, &recvTime);
        waitTime = GetMicroSecond(recvTime);
        if (waitTime < m_nextSendTime) {
            waitTime = m_nextSendTime - waitTime;
        } else {
            waitTime = 0;
        }

        // avoid waitTime is 0 to m_packetQueue.GetNextPktWait
        if (waitTime < GET_DATA_FROM_QUEUE_MIN_WAIT_TIME) {
            waitTime = GET_DATA_FROM_QUEUE_MIN_WAIT_TIME;
        }

        recvMessage = reinterpret_cast<StreamMsgHead *>(m_packetQueue.GetNextPktWait(
            waitTime / SECOND_TO_MILL));
        if (recvMessage == nullptr) {
            break;
        }

        clock_gettime(CLOCK_MONOTONIC, &recvTime);
        responseMessage = m_msgReassemble.Reassemble(*recvMessage);
        uint8_t* msgData = responseMessage.first;
        uint32_t msgSize = responseMessage.second;
        if (msgData == nullptr) {
            break;
        }

        if (msgSize >= sizeof(HeartbeatMsg)) {
            HandleRecvMsg(responseMessage, recvTime);
        } else {
            ERR("Heartbeat response msg size(%u) is error", msgSize);
        }

        VmiFreeBuffer(msgData);
        responseMessage = std::pair<uint8_t*, uint32_t>(nullptr, 0);
    } while (recvMessage != nullptr);
}

/**
 * @brief: compute net loop
 * @param [in] heartbeatResponseMsg: heartbeat package
 * @param [in] recvTime: recvTime
 */
void Heartbeat::ThreadPara::CalcLatency(HeartbeatMsg& heartbeatResponseMsg, uint64_t recvTime)
{
    if (heartbeatResponseMsg.seqNum != m_recvSeqNum + 1) {
        ERR("Heartbeat response msg order(%ju) is error, expected value is %ju." \
            " There may be a packet loss on the network.!",
            heartbeatResponseMsg.seqNum, m_recvSeqNum + 1);
    }

    m_recvSeqNum = heartbeatResponseMsg.seqNum;
    uint64_t latencyMicroSec = recvTime - heartbeatResponseMsg.sendMicroSecond;
    int size = m_latencyQueue.size();

    for (int i = size - 1; i >= 0; --i) {
        if (m_latencyQueue.at(i).first == heartbeatResponseMsg.sendMicroSecond) {
            m_latencyQueue.at(i).second = latencyMicroSec;
            DBG("heartbeat msg: seq = %ju, send_time = %ju, latency = %ju",
                heartbeatResponseMsg.seqNum,
                heartbeatResponseMsg.sendMicroSecond,
                latencyMicroSec);
            break;
        }
    }  // If not found, it may be a network exception and does not process the data.

    latencyMicroSec = 0;
    uint64_t maxMicroSec = 0;
    int receivePacketCount = 0;
    for (int i = size - 1; (i >= 0) && (i > size - 1 - AVERAGE_NUM); --i) {
        uint64_t tmp = m_latencyQueue.at(i).second;
        latencyMicroSec += tmp;
        maxMicroSec = (maxMicroSec < tmp) ? tmp : maxMicroSec;
        receivePacketCount += (tmp == 0) ? 0 : 1;
    }

    receivePacketCount = (receivePacketCount > 0) ? receivePacketCount : 1;

    m_netLoopbackLatencyMax = maxMicroSec;
    m_netLoopbackLatencyAverage = latencyMicroSec / receivePacketCount;
    LOG_RATE_LIMIT(ANDROID_LOG_INFO, 1, "MaxLatency = %ju, AverageLatency = %ju",
                   maxMicroSec, m_netLoopbackLatencyAverage);

#if ENABLE_STATISTICS
    // statistics heartbeat
    Statistics* statistics = Statistics::GetInstance();
    if (statistics != nullptr) {
        static StatisticsObject *netStatisticsObject = statistics->AllocStatisticsObject(NET_TABLE_NAME,
                                                                                         NET_TABLE_COLUMNS,
                                                                                         SAMPLE_INTERVAL_PERSECOND);
        if (netStatisticsObject != nullptr) {
            static int indexHeartLoopBackTime = netStatisticsObject->GetComponentIndex(NET_COLUMN_LOOP_TIME);
            netStatisticsObject->UpdateComponent(indexHeartLoopBackTime, UPDATE, m_netLoopbackLatencyAverage);
        }
    }
#endif

    return;
}

/**
 * @brief: check net status
 */
void Heartbeat::ThreadPara::CheckNetStatus()
{
    int queueSize = m_latencyQueue.size();
    uint64_t receivePacketCount = 0;

    for (int i = queueSize - 1; i >= 0; --i) {
        uint64_t tmp = m_latencyQueue.at(i).second;
        receivePacketCount += (tmp == 0) ? 0 : 1;
    }

    if ((queueSize >= CHECK_DURATION) && (receivePacketCount == 0)) {
        ERR("Network may be interrupted or delayed too much.");

        m_stopFlag = true;
        if (m_callback != nullptr) {
            std::thread t(
                [this]() -> void {
                    INFO("execute m_callback");
                    m_callback(m_vmiSocket);
                    return;
                });//lint !e527
            if (t.joinable()) {
                t.detach();
            }
        } else {
            ERR("m_callback is nullptr!");
        }
    }

    return;
}

/**
 * @brirf: construct
 * @param [in] sock: socket
 * @param [in] fnCallback: check net exception, callbak it
 */
Heartbeat::Heartbeat(std::shared_ptr<VmiSocket> socket, std::function<void(std::shared_ptr<VmiSocket>)> fnCallback)
    : m_lock()
{
    m_vmiSocket = socket;
    m_callback = fnCallback;
    m_task = nullptr;
}

Heartbeat::~Heartbeat()
{
}

/**
 * @brief: start thread
 */
bool Heartbeat::Start()
{
    std::lock_guard<std::mutex> lck(m_lock);  // Prevent multiple startups
    if (m_threadParameter == nullptr) {
        m_threadParameter = std::shared_ptr<Heartbeat::ThreadPara>(new (std::nothrow) Heartbeat::ThreadPara());
        if (m_threadParameter == nullptr) {
            ERR("error: failed to alloc thread parameter");
            return false;
        }

        m_threadParameter->SetVmiSocket(m_vmiSocket);
        m_threadParameter->SetCallback(m_callback);
        m_threadParameter->SetStopFlag(false);

        m_task = new (std::nothrow) std::thread (TaskentryHeartbeat, m_threadParameter);
        if (m_task == nullptr) {
            ERR("error: failed to alloc task");
            m_threadParameter = nullptr;
            return false;
        }
    }

    return true;
}

/**
 * @brief: Stop the thread and the caller waits for the thread to exit and hear.
 */
void Heartbeat::ForceStop()
{
    std::lock_guard<std::mutex> lck(m_lock);

    if (m_threadParameter != nullptr) {
        m_threadParameter->SetStopFlag(true);
        m_threadParameter->GetInterThreadComm().Cancel();
        if (m_task != nullptr) {
            m_task->join();
            delete m_task;
            m_task = nullptr;
        }
        ERR("heartbeat thread exit\n");
        m_threadParameter = nullptr;
    }

    return;
}

/**
 * @brief: put heartbeat data into queue
 * @param [in] packet: heartbeat data
 */
void Heartbeat::Handle(const VmiBuffer& packet)
{
    std::lock_guard<std::mutex> lck(m_lock);
    if (m_threadParameter != nullptr) {
        m_threadParameter->GetInterThreadComm().PutPkt(packet);
    } else {
        VmiFreeBuffer(packet.GetPointer());
    }
}

/**
 * @brief: Get network average latency
 * @return: latency, 0 is Infinity
 */
uint64_t Heartbeat::GetNetLoopbackLatencyAverage()
{
    if (m_threadParameter != nullptr) {
        return m_threadParameter->GetNetLoopbackLatencyAverage();
    }

    return 0;
}

/**
 * @brief: Get the maximum network latency
 * @return: latency, 0 is Infinity
 */
uint64_t Heartbeat::GetNetLoopbackLatencyMax()
{
    if (m_threadParameter != nullptr) {
        return m_threadParameter->GetNetLoopbackLatencyMax();
    }

    return 0;
}
